"""
particle swarm optimization for 4-dim function
"""
import random
import sys
import matplotlib.pyplot as plt

DIM = 4
LB = -30  # lower bound
UB = 30  # upper bound
V_UB = 0.15 * (UB-LB)  # upper bound of velocity
W_MAX = 0.9
W_MIN = 0.4
C1 = 2
C2 = 2
# bigger the w, more persistence in self_vel
ITER_MAX = 300
POPU_SIZE = 100


class particle():
    def __init__(self):
        self.pos = [random.uniform(LB, UB) for i in range(DIM)]
        self.v = [random.uniform(-V_UB, V_UB) for i in range(DIM)]
        self.best = self.pos

    def update(self, w, g_best_slu):
        old_pos = self.pos  # hold the old pos value
        old_v = self.v
        rand1 = random.random()
        rand2 = random.random()
        delta_self = [self.best[i] - self.pos[i] for i in range(DIM)]
        delta_global = [g_best_slu[i] - self.pos[i] for i in range(DIM)]
        for i in range(DIM):
            self.v[i] = w * self.v[i] + C1 * rand1 * \
                delta_self[i] + C2 * rand2 * delta_global[i]
            if (self.v[i] > V_UB):
                self.v[i] = V_UB
            if (self.v[i] < -V_UB):
                self.v[i] = -V_UB
        for i in range(DIM):
            self.pos[i] = old_pos[i] + old_v[i]
            if (self.pos[i] > UB):
                self.pos[i] = UB
            if (self.pos[i] < LB):
                self.pos[i] = LB
        if f(self.pos) < f(self.best):
            self.best = self.pos


def avg_swarm(swarm):
    avg = [0 for i in range(DIM)]
    for i in range(DIM):
        total = 0
        for p in swarm:
            total = total + p.pos[i]
        avg[i] = total/len(swarm)
    return avg


def avg_swarm_v(swarm):
    avg = [0 for i in range(DIM)]
    for i in range(DIM):
        total = 0
        for p in swarm:
            total = total + p.v[i]
        avg[i] = total/len(swarm)
    return avg


def update_w(w, n_iter):
    w = W_MAX - (W_MAX - W_MIN) * (n_iter / ITER_MAX)
    return w


def f(x):
    s = 0
    for i in range(DIM-1):
        s = s + 100*(x[i+1]-(x[i])**2)**2 + (x[i]-1)**2
    return s


def sel_best(swarm):
    best_val = 1e16
    best_sol = 0
    for p in swarm:
        val = f(p.pos)
        if val < best_val:
            best_sol = p
            best_val = val
    best_sol = best_sol.pos[:]
    return best_sol, best_val


conn = [1, 2, 3, 4]


def pso_func(f):

    x = []
    y = []
    plt.xlabel("ITER")
    plt.ylabel("F")
    swarm = [particle() for i in range(POPU_SIZE)]
    g_best_slu, g_best_val = sel_best(swarm)

    # print("WEIRD:", g_best.pos)
    # F00 = f(g_best.pos)
    # if F00 != g_best_val:
    #     print("OUTER ERROR")

    w = W_MAX
    for n_iter in range(ITER_MAX):
        w = update_w(w, n_iter)

        for p in swarm:
            p.update(w, g_best_slu)
        cur_best_slu, cur_best_val = sel_best(swarm)

        # print("WEIRD:", g_best.pos)
        # F0 = f(g_best.pos)
        # if F0 != g_best_val:
        #     print("OUTER ERROR")

        # if f(cur_best.pos) == cur_best_val:
        #     print("MATCH")

        if cur_best_val < g_best_val:
            g_best_slu = cur_best_slu
            g_best_val = cur_best_val

            # F1 = f(g_best.pos)
            # if F1 != g_best_val:
            #     print("INNER ERROR")

        # F2 = f(g_best.pos)
        # if F2 != g_best_val:
        #     print("ERROR")

        print("【ITER:", n_iter, "】")
        print("\tglobal best:", g_best_slu)
        avg = avg_swarm(swarm)
        avg_delta = [avg[i] - g_best_slu[i] for i in range(DIM)]
        print("\tavg - global:", avg_delta)
        print("\tavg v", avg_swarm_v(swarm))
        print("\tw:", w)
        print("\tF:", f(g_best_slu))
        if g_best_val < 10:
            x.append(n_iter)
            y.append(g_best_val)
        # print("\tG:", g_best_val)
    plt.plot(x, y)
    plt.show()


with open('E:/codes/courses/ArtificialIntelligence/PSO/out_dynamic_w.txt', 'w+', encoding='utf-8') as file:
    sys.stdout = file
    pso_func(f)
